# Import libraries

from Func_Extract_Data import Extract_Track_Data
from Func_Create_File import create_file
from Func_List_Files import list_ncfiles
from netCDF4 import Dataset
from wrf import getvar
import numpy as np
import csv
import os

# Define settings
#----------------------------------------------------------------------------------------------
Input_Dir  = '.../Hurricane_Case/'           # Provide the input directory   
Output_Dir = '.../Post_Processing/Track/'    # Provide the output directory   
HNS        = ['Florence']
HRS        = ['96']
GSS        = ['8km']
TMS        = ['NoTurb']
CLTS       = ['TSKM2', 'Default', 'TSKP2']
radius     = 80

# Extract the track
#----------------------------------------------------------------------------------------------
Idx = 0
for HN in HNS:
    for GS in GSS:
        for TM in TMS:
            clt_count = 0
            for CLT in CLTS:
                Real_Lats, Real_Lons, ncfiles = [], [], []
                min_slp, min_lat, min_lon     = [], [], []

                Hurricane_Setting = HN + '_' + HRS[Idx] + 'Hours_' + GS + '_' + TM + '_' + CLTS[clt_count]
                Input_Dir_1       = Input_Dir + Hurricane_Setting
                Real_Lats         = Extract_Track_Data (Output_Dir, 'Real_Output.csv', Real_Lats, 'Real_lat') # These latitudes will be used to locate the first location
                Real_Lons         = Extract_Track_Data (Output_Dir, 'Real_Output.csv', Real_Lons, 'Real_long') # These longitudes will be used to locate the first location
                ncfiles           = list_ncfiles (Input_Dir_1, ncfiles)
                
                for ncfile in ncfiles:
                    os.chdir(Input_Dir_1)
                    Data     = Dataset(ncfile)
                    lat      = np.array(getvar(Data, "XLAT"))[:,0]
                    lon      = np.array(getvar(Data, "XLONG"))[0,:]
                    slp      = np.array(getvar(Data, "slp", 0))
                    lat_idx  = 0
                    lon_idx  = 0

                    if ncfile == ncfiles[0]:
                        for i in range(len(lat)-1):
                            if lat[i] < Real_Lats[0] and lat[i+1] > Real_Lats[0]:
                                lat_idx = i
                        for j in range(len(lon)-1):
                            if lon[j] < Real_Lons[0] and lon[j+1] > Real_Lons[0]:
                                lon_idx = j
                        
                        i_idx = min(radius, abs(len(lat)-lat_idx), lat_idx)
                        j_idx = min(radius, abs(len(lon)-lon_idx), lon_idx)
                        Eye   = np.where(slp == np.amin(slp[lat_idx-i_idx:lat_idx+i_idx, lon_idx-j_idx:lon_idx+j_idx]))
                        min_lat.append(float(lat[Eye[0][0]]))
                        min_lon.append(float(lon[Eye[1][0]]))

                    else:
                        lat_idx = int(np.where(lat == min_lat[-1:])[0])
                        lon_idx = int(np.where(lon == min_lon[-1:])[0])

                        i_idx = min(radius, abs(len(lat)-lat_idx), lat_idx)
                        j_idx = min(radius, abs(len(lon)-lon_idx), lon_idx)
                        Eye   = np.where(slp[lat_idx-i_idx:lat_idx+i_idx, lon_idx-j_idx:lon_idx+j_idx] == np.amin(slp[lat_idx-i_idx:lat_idx+i_idx, lon_idx-j_idx:lon_idx+j_idx]))
                        min_lat.append(float(lat[Eye[0][0] + lat_idx-i_idx]))
                        min_lon.append(float(lon[Eye[1][0] + lon_idx-j_idx]))
# Export the output
#----------------------------------------------------------------------------------------------
                os.chdir(Output_Dir)
                Path = os.getcwd()
                MyFile=open('%s.csv' %Hurricane_Setting,'w')
                MyFile.write (Hurricane_Setting + "\n")
                MyFile.write ("min_lat,min_long\n")
                for i in range(len(min_lat)):
                        MyFile.write (str(min_lat[i]) + ',' + str(min_lon[i]) + "\n")
                MyFile.close()                                 
                clt_count += 1
    Idx += 1
